import os
import glob
import numpy as np


from daniel_tools.img_utils import colorize_depth_map,concat_images, imgs2video2


import pyexr
import json
from os.path import join
import cv2

from loguru import logger


def load_jsonl(jsonl_path):
    with open(jsonl_path, 'r') as f:
        return [json.loads(line) for line in f]


def load_data_by_seq_name(seq_name):
    for line in gt_data:
        if line['seq_name'] == seq_name:
            return line
    return None



def load_depth(depth_paths, depth_range):

    
    
    depth_list = []
    mask_list = []

    for dp in depth_paths:
        depth = cv2.imread(dp, -1)
        depth = depth.astype(np.float32) / 1000.0
        mask = (depth > depth_range[0]) & (depth < depth_range[1])
        depth_list.append(depth)
        mask_list.append(mask)

    depth_list = np.stack(depth_list, axis=0)
    mask_list = np.stack(mask_list, axis=0)



    return depth_list,mask_list



def load_depth_exr(depth_paths, depth_range):
    

    depth_list = []
    mask_list = []

    for dp in depth_paths:
        depth = pyexr.open(dp).get()
        depth = depth.squeeze()
        
        
        
        mask = (depth > depth_range[0]) & (depth < depth_range[1])
        depth_list.append(depth)
        mask_list.append(mask)

    depth_list = np.stack(depth_list, axis=0)
    mask_list = np.stack(mask_list, axis=0)



    return depth_list,mask_list

# gt_jsonl_path = "data/pics/pics.jsonl"

# gt_jsonl_path = "data/phocal/test_all.jsonl"
# prediction_path = 'workspace/phocal/test_all.jsonl--1-448-640'


import argparse




def resize_3d_array(depth, new_width, new_height):
    """
    Resize a (F, H, W) array and its mask to (F, new_height, new_width).
    """
    F = depth.shape[0]
    resized_depth = np.zeros((F, new_height, new_width), dtype=depth.dtype)
    for i in range(F):
        resized_depth[i] = cv2.resize(depth[i], (new_width, new_height), interpolation=cv2.INTER_NEAREST)
    return resized_depth



parser = argparse.ArgumentParser()
parser.add_argument('--gt_jsonl_path', type=str, default='data/clearpose/test.jsonl')
parser.add_argument('--prediction_path', type=str, default='workspace/clearpose/test.jsonl-400')

args = parser.parse_args()

gt_jsonl_path = args.gt_jsonl_path
prediction_path = args.prediction_path


gt_data = load_jsonl(gt_jsonl_path)
gt_root = os.path.dirname(gt_jsonl_path)


from PIL import Image

from tqdm import tqdm

for npz_path in tqdm(glob.glob(f"{prediction_path}/**/*.npz"),desc='out side loop'):
    seq_name = npz_path.split('/')[-2]

    
    gt_seq_data = load_data_by_seq_name(seq_name)


    #* load rgb 
    if gt_seq_data.get('rgb-suffix', None) is not None:
        gt_all_imgs = sorted(glob.glob(join(gt_root,gt_seq_data['rgb'], f'*{gt_seq_data["rgb-suffix"]}')))
    else:
        gt_all_imgs = sorted(glob.glob(join(gt_root,gt_seq_data['rgb'], '*.png')))

    
    #* load gt depth
    
    if gt_seq_data.get('depth-suffix', None) is not None:
        all_depth_paths = sorted(glob.glob(join(gt_root,gt_seq_data['rgb'],f'*{gt_seq_data["depth-suffix"]}')))
    else:
        all_depth_paths = sorted(glob.glob(join(gt_root,gt_seq_data['depth'],'*')))

    
    load_function =load_depth_exr if all_depth_paths[0].endswith('.exr') else load_depth

    breakpoint()
    if gt_seq_data.get('depth-range', None) is not None:
        gt_depth,gt_mask = load_function(all_depth_paths, gt_seq_data['depth-range'])
    else:
        logger.info(f'{seq_name} has no depth-range, use default depth-range [1e-3,100000]')
        gt_depth,gt_mask = load_function(all_depth_paths, [1e-3,5])
        if gt_depth.shape[-1] ==4 :
            gt_depth = gt_depth[...,:3].mean(axis=-1)
            gt_mask = gt_mask[...,:3].mean(axis=-1)
            gt_mask = gt_mask.astype(np.bool_)
        elif gt_depth.shape[-1]  == 1:
            gt_depth = gt_depth.squeeze(-1)
            gt_mask = gt_mask.squeeze(-1)

            
    with np.load(npz_path) as data:
        all_pred_depths = data['depth'].astype(np.float32).squeeze()


    F1,H1,W1 = all_pred_depths.shape
    F2, H2, W2 = gt_depth.shape

    if H1 != H2 or W2 != W1:
        all_pred_depths = resize_3d_array(all_pred_depths,  W2,H2)
        logger.info(f'after resizing, {all_pred_depths.shape}')

    
    imgs = []
    


    prediction_save_dir = join(os.path.dirname(npz_path),'predictions')
    os.makedirs(prediction_save_dir,exist_ok=True)


    for rgb_path,x,gt_depth,gt_mask in tqdm(zip(gt_all_imgs,all_pred_depths,gt_depth,gt_mask),desc=f'{seq_name}'):
        img_idx = os.path.basename(rgb_path).split('.')[0]
        # rgb_img = Image.open(rgb_path)
        
        # gt_depth[~gt_mask] = 1e-3
        x[~gt_mask] = 1e-3

        x =colorize_depth_map(x, mask=gt_mask)

        # gt_depth = colorize_depth_map(gt_depth)
        # gt_mask = colorize_depth_map(gt_mask.astype(np.uint8))
        # merged_img = concat_images([rgb_img,x,gt_depth,gt_mask])
        
        x.save(join(prediction_save_dir,f'{img_idx}.jpg'))


    
    # imgs2video2(imgs, out_path = npz_path.replace('.npz', 'concated.mp4'), fps=10)
    print(npz_path,'is done')
    
    